use std::iter::once;

use num_integer::Integer;
use rand::{distributions::Uniform, prelude::Distribution, thread_rng, Rng};

/// Safe modular arithmetic without overflow.
pub trait ModComputation {
    /// Calculates `(self + other) % modulus` without overflow.
    fn add_mod(self, other: Self, modulus: Self) -> Self;
    /// Calculates `(self << other) % modulus` without overflow.
    fn shl_mod(self, other: u32, modulus: Self) -> Self;
    /// Calculates `self * other % modulus` without overflow.
    fn mul_mod(self, other: Self, modulus: Self) -> Self;
    /// Calculates `self ^ exp % modulus` efficiently, using the "square and multiply" method.
    fn pow_mod(self, exp: Self, modulus: Self) -> Self;
}

impl ModComputation for u128 {
    fn add_mod(mut self, mut other: Self, modulus: Self) -> Self {
        self %= modulus;
        other %= modulus;

        let (sum, overflow) = self.overflowing_add(other);
        if !overflow {
            sum % modulus
        } else {
            // `modulus` is large, but `self + other` is smaller than `2 * modulus`.
            sum.wrapping_sub(modulus)
        }
    }

    fn shl_mod(mut self, mut other: u32, modulus: Self) -> Self {
        self %= modulus;
        if self == 0 {
            return 0;
        }

        while other != 0 {
            let shift = self.leading_zeros().min(other);
            if shift == 0 {
                self = self.wrapping_shl(1).wrapping_sub(modulus);
                other -= 1;
            } else {
                self = (self << shift) % modulus;
                other -= shift;
            }
        }
        self
    }

    fn mul_mod(mut self, mut other: Self, modulus: Self) -> Self {
        // For small `modulus` or small `self * other`, we can use the naive method.
        self %= modulus;
        other %= modulus;
        if let Some(product) = self.checked_mul(other) {
            product % modulus
        } else {
            let a_hi = self >> 64;
            let a_lo = self & 0xFFFF_FFFF_FFFF_FFFF;
            let b_hi = other >> 64;
            let b_lo = other & 0xFFFF_FFFF_FFFF_FFFF;

            (a_hi * b_hi)
                .shl_mod(128, modulus)
                .add_mod(
                    (a_hi * b_lo)
                        .add_mod(a_lo * b_hi, modulus)
                        .shl_mod(64, modulus),
                    modulus,
                )
                .add_mod(a_lo * b_lo, modulus)
        }
    }

    fn pow_mod(mut self, mut exp: Self, modulus: Self) -> Self {
        if self == 1 || exp == 0 {
            return 1;
        }
        let mut acc = 1;

        loop {
            if (exp & 1) == 1 {
                acc = acc.mul_mod(self, modulus);
                if exp == 1 {
                    return acc;
                }
            }
            exp >>= 1;
            self = self.mul_mod(self, modulus);
        }
    }
}

const SMALL_ODD_PRIMES: &[u128] = &[
    3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97,
];

/// A wrapper function for `miller_rabin`, to test whether a number is prime.
///
/// `n` should be checked against small primes
fn is_prime_mr(n: u128) -> bool {
    // Miller-Rabin bases selected from <https://miller-rabin.appspot.com/>.
    // over 2^18
    if n < 341531 {
        miller_rabin(n, once(9345883071009581737))
    }
    // over 2^29
    else if n < 1050535501 {
        miller_rabin(n, [336781006125, 9639812373923155])
    }
    // over 2^38
    else if n < 350269456337 {
        miller_rabin(
            n,
            [
                4230279247111683200,
                14694767155120705706,
                16641139526367750375,
            ],
        )
    }
    // over 2^38
    else if n < 55245642489451 {
        miller_rabin(
            n,
            [
                2,
                141889084524735,
                1199124725622454117,
                11096072698276303650,
            ],
        )
    }
    // over 2^52
    else if n < 7999252175582851 {
        miller_rabin(
            n,
            [
                2,
                4130806001517,
                149795463772692060,
                186635894390467037,
                3967304179347715805,
            ],
        )
    }
    // over 2^52
    else if n < 585226005592931977 {
        miller_rabin(
            n,
            [
                2,
                123635709730000,
                9233062284813009,
                43835965440333360,
                761179012939631437,
                1263739024124850375,
            ],
        )
    } else if n < 1 << 64 {
        miller_rabin(n, [2, 325, 9375, 28178, 450775, 9780504, 1795265022])
    }
    // OEIS A014233: https://oeis.org/A014233
    // over 2^78
    else if n < 318665857834031151167461 {
        miller_rabin(n, [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37])
    }
    // over 2^81
    else if n < 3317044064679887385961981 {
        miller_rabin(n, [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41])
    } else {
        // No known definite bases for `number`. Use a random base sequence with 10 bases.
        // The probability of a false positive is less than 4^-10, less than one in a million.
        miller_rabin(n, thread_rng().sample_iter(Uniform::new(2, n - 1)).take(10))
    }
}

/// Performs the Miller-Rabin primality test on `n` using the given `bases`.
///
/// If `n` is composite for any `base`, then `false` is returned.
/// Else `true` is returned, indicating that `n` is probably prime.
fn miller_rabin<I>(n: u128, bases: I) -> bool
where
    I: IntoIterator<Item = u128>,
{
    for base in bases {
        // Miller-Rabin test is based on the following theorem:
        //     If n is an odd prime,
        //     let n - 1 = d * 2^s, where d is odd.
        //     Then ∀a coprime to n,
        //         a^d ≡ 1 (mod n)                        (1), or
        //         ∃r ∈ [0, s), a^(2^r * d) ≡ -1 (mod n)  (2)
        let mut d = n - 1;
        let s = d.trailing_zeros();
        d >>= s;

        let mut x = base.pow_mod(d, n);

        if x == 1 {
            continue;
        }

        for r in 0..s {
            if x == n - 1 {
                break;
            } else if r == s - 1 {
                // We failed to find an `r` as described in (2).
                return false;
            }

            x = x.mul_mod(x, n);

            if x == 1 {
                // A non-trivial square root of 1 modulo `n` is found.
                // If `n` is prime, the only possible solutions of x^2 ≡ 1 (mod n)
                // are x ≡ 1 (mod n) and x ≡ 1 (mod n).
                // `n` must be composite.
                return false;
            }
        }
    }
    // All tests has passed, `n` is a probably prime, given the bases.
    true
}

/// Find a non-trivial factor of the given num.
///
/// A non-trivial factor is a factor that is neither 1 nor the number itself.
///
/// The current implementation is based on Pollard's Rho Algorithm.
/// For a number n with minimal prime factor p, the expected time is O(p ^ (1 / 2)),
/// which is less than O(n ^ (1 / 4)) if n is a composite number.
///
/// **WARNING**: A prime input may cause an infinite loop.
fn non_trivial_factor(n: u128) -> u128 {
    macro_rules! return_if_non_trivial {
        ($e:expr) => {
            if let d @ 2.. = $e {
                return d;
            }
        };
    }

    let mut rng = thread_rng();
    let x0_range = Uniform::new(0, n);
    let c_range = Uniform::new(1, n - 2);

    let init_f = |c| move |x: u128| x.mul_mod(x, n).add_mod(c, n);

    let mut c = c_range.sample(&mut rng);
    let mut f = init_f(c);
    let mut x0 = x0_range.sample(&mut rng);

    loop {
        let mut tortoise = x0;
        let mut hare = f(x0);
        let mut power = 1;
        let mut product = 1;

        'a: loop {
            for lambda in 0..power {
                hare = f(hare);
                product = product.mul_mod(hare.abs_diff(tortoise), n);

                if product == 0 {
                    // We got a trivial factor n, try another (c, x0) pair.
                    break 'a;
                }

                if lambda % 127 == 0 {
                    return_if_non_trivial!(product.gcd(&n));
                }
            }

            return_if_non_trivial!(product.gcd(&n));

            tortoise = hare;
            power <<= 1;
            product = 1;
        }
        x0 = x0_range.sample(&mut rng);
        c = c_range.sample(&mut rng);
        f = init_f(c);
    }
}

pub fn find_max_prime_factor(mut number: u128) -> u128 {
    let mut max_prime_factor = 1;

    if (number & 1) == 0 {
        max_prime_factor = 2;
        number >>= number.trailing_zeros();
    }

    // Trial division
    for &prime in SMALL_ODD_PRIMES {
        if number % prime == 0 {
            max_prime_factor = prime;
            loop {
                number /= prime;
                if number % prime != 0 {
                    break;
                }
            }
        }

        if number == 1 {
            return max_prime_factor;
        }
    }

    if is_prime_mr(number) {
        // number is a prime larger than any one in `SMALL_PRIMES`
        return number;
    }

    let mut composite_factors = vec![number];
    while let Some(mut n) = composite_factors.pop() {
        let f = non_trivial_factor(n);

        loop {
            // Factor out f from n.
            n /= f;
            if n % f != 0 {
                break;
            }
        }

        if is_prime_mr(f) {
            max_prime_factor = max_prime_factor.max(f);
        } else {
            composite_factors.push(f);
        }

        if is_prime_mr(n) {
            max_prime_factor = max_prime_factor.max(n);
        } else {
            composite_factors.push(n);
        }
    }

    max_prime_factor
}
